from ssd import build_ssd
from data import *
import torch
import cv2
import numpy as np
from torch.autograd import Variable
from matplotlib import pyplot as plt
from xml.dom.minidom import parse
import os.path as osp
import os
from collections import Counter
from utils import Filter
import shutil

# tongzhou20190831 tongzhou20190901 tongzhou20190902 tongzhou20190904 tongzhou20191001
root_path = '/home/ubuntu/code/ssd.pytorch/data/pollen_data/test'
images_path = osp.join(root_path, 'JPEGImages')

# load net
if torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
num_classes = len(VOC_CLASSES) + 1  # +1 background
net = build_ssd('test', 300, num_classes)  # initialize SSD
net.load_weights('./weights/VOC_2class_1008.pth')
net.to(device='cuda')

# 图片过滤
lower_red = np.array((54, 22, 82))
upper_red = np.array((255, 250, 250))

f = Filter()
CENTERS = np.array([[219.57673667, 220.39418417, 228.82714055],
                    [219.89855072, 181.42028986, 243.42028986]])


def predict(X: np.ndarray = None, centers: np.ndarray = CENTERS):
    if X is None:
        return None
    x = np.broadcast_to(X, centers.shape)
    y = np.sum((centers - x) ** 2, axis=1)
    return y.argmin(axis=0)


def filter(image, cls):
    im = image.copy()
    contrast = im * 1.5
    contrast[contrast > 255] = 255
    contrast = contrast.astype(np.uint8)
    if cls == 0:
        lower_pink = np.array([150, 20, 20])

        upper_pink = np.array([179, 255, 255])
    else:
        lower_pink = np.array([150, 80, 80])

        upper_pink = np.array([179, 255, 255])
    hsv_image = cv2.cvtColor(contrast, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv_image, lower_pink, upper_pink)
    closing = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((30, 30), np.uint8))
    # closing = cv2.morphologyEx(closing, cv2.MORPH_CLOSE, np.ones((10, 10), np.uint8))
    # closing = cv2.morphologyEx(closing, cv2.MORPH_CLOSE, np.ones((10, 10), np.uint8))
    opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, np.ones((10, 10), np.uint8))
    # after_filted = cv2.bitwise_and(image, image, mask=opening)
    cnts, _ = cv2.findContours(opening, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # cv2.imwrite('opne.jpg', opening)
    # cv2.imwrite('mask.jpg', mask)
    return cnts, opening


def image_filter(image: np.ndarray) -> np.ndarray:
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, lower_red, upper_red)
    return cv2.bitwise_and(image, image, mask=mask)


def predict_new(image_path, save=True, score_threshold=0.8, default_means=True, dir=None):
    image = cv2.imread(image_path)
    if image is None:
        return
    classification = predict(np.expand_dims(np.percentile(image, 50, axis=(0, 1)), 0))
    _, dst = filter(image, classification)
    contours, hierarchy = cv2.findContours(dst, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) == 0:
        return
    bg_image = np.broadcast_to(config.MEANS, image.shape).astype('uint8')
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        xmin = x - 10 if x - 10 > 0 else 0
        ymin = y - 10 if y - 10 > 0 else 0
        xmax = xmin + w + 15
        ymax = ymin + h + 15
        bg_image[ymin:ymax, xmin:xmax, :] = image[ymin:ymax, xmin:xmax, :]

    # cv2.imwrite('hhh.jpg', image)
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.figure(figsize=(10, 10))

    # lower_blue = np.array((100, 40, 20))
    # upper_blue = np.array((255, 230, 236))
    # hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
    # mask = cv2.inRange(hsv, lower_blue, upper_blue)
    # im = cv2.bitwise_and(image, image, mask=mask)

    if default_means:
        x = cv2.resize(bg_image, (300, 300)).astype(np.float32)
        x -= MEANS  # MEANS
    else:
        im = image.astype(np.float32) * 1.1
        im[im > 255] = 255
        # im = image_filter(image)
        x = cv2.resize(im, (300, 300)).astype(np.float32)
        x -= MEANS
    x = x.astype(np.float32)
    x = x[:, :, ::-1].copy()

    x = torch.from_numpy(x).permute(2, 0, 1)

    with torch.no_grad():

        xx = Variable(x.unsqueeze(0))  # wrap tensor in Variable
        if torch.cuda.is_available():
            xx = xx.cuda()

        y = net(xx)

        from data import VOC_CLASSES as labels

        top_k = 10

        colors = plt.cm.hsv(np.linspace(0, 1, num_classes)).tolist()
        plt.imshow(rgb_image)  # plot the image for matplotlib
        currentAxis = plt.gca()

        detections = y.data
        # scale each detection back up to the image
        scale = torch.tensor(rgb_image.shape[1::-1], dtype=torch.float32).repeat(2)
        # for i in detections[0]:
        #     print(i)
        # return
        have_targets = False
        for i in range(detections.size(1)):
            j = 0
            if i == 0:
                continue
            while detections[0, i, j, 0] >= score_threshold:
                score = detections[0, i, j, 0]
                label_name = labels[i - 1]
                display_txt = '%s: %.2f' % (label_name, score)
                pt = (detections[0, i, j, 1:] * scale).cpu().numpy()
                coords = (pt[0], pt[1]), pt[2] - pt[0] + 1, pt[3] - pt[1] + 1
                color = colors[i]
                currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
                currentAxis.text(pt[0], pt[1], display_txt, bbox={'facecolor': color, 'alpha': 0.5})
                j += 1
                have_targets = True
        if save and have_targets:
            if not osp.exists(osp.join('/home/ubuntu/z7z8/detection/1026', dir)):
                os.mkdir(osp.join('/home/ubuntu/z7z8/detection/1026', dir))
            plt.savefig(osp.join('/home/ubuntu/z7z8/detection/1026', dir, image_path.split("/")[-1]))
            if not osp.exists(osp.join('/home/ubuntu/z7z8/result/1026', dir)):
                os.mkdir(osp.join('/home/ubuntu/z7z8/result/1026', dir))
            shutil.copy(image_path, osp.join('/home/ubuntu/z7z8/result/1026', dir))
        # plt.show()
        plt.cla()
        plt.close('all')

def run(arr):
    for dir in arr:
        images_path = osp.join('/media/ubuntu/Extreme SSD/data10.26', dir)
        for i in os.listdir(images_path):
            image_path = osp.join(images_path, i)
            print(image_path)
            predict_new(image_path, True, 0.8, default_means=True, dir=dir)

if __name__ == '__main__':
    from multiprocessing import Process
    import multiprocessing
    arr = os.listdir('/media/ubuntu/Extreme SSD/data10.26')
    process_num = 6
    print(process_num)


    processes = []
    if len(arr) < process_num:
        num = 1
        multiprocessing.set_start_method('spawn')
        for i in range(process_num):
            processes.append(Process(target=run, args=(arr[i * process_num:(i + 1) * process_num],)))
        for p in processes:
            p.start()
        for p in processes:
            p.join()
    else:
        multiprocessing.set_start_method('spawn')
        num = len(arr) // process_num
        for i in range(process_num):
            if i < process_num - 1:
                processes.append(Process(target=run, args=(arr[i * num:(i + 1) * num],)))
            else:
                processes.append(Process(target=run, args=(arr[i * num:],)))
        print(len(processes))
        for p in processes:
            p.start()
        for p in processes:
            p.join()
